import torch
import torchvision
import torchvision.transforms as transforms
from torch import nn
from torchvision.models import vit_b_16, ViT_B_16_Weights


# 1. 数据准备
def get_cifar10_data():
    # ViT预训练模型期望的输入大小是224x224
    transform = transforms.Compose([
        transforms.Resize(224),  # 调整图像大小到224x224
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])  # ImageNet标准化参数
    ])

    trainset = torchvision.datasets.CIFAR10(
        root='./data', train=True, download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=64, shuffle=True, num_workers=2)

    testset = torchvision.datasets.CIFAR10(
        root='./data', train=False, download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(
        testset, batch_size=64, shuffle=False, num_workers=2)

    classes = ('plane', 'car', 'bird', 'cat', 'deer',
               'dog', 'frog', 'horse', 'ship', 'truck')

    return trainloader, testloader, classes


# 2. 修改ViT模型
class ViTForCIFAR10(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        # 加载预训练模型
        self.vit = vit_b_16(weights=ViT_B_16_Weights.IMAGENET1K_V1)

        # 冻结所有参数
        for param in self.vit.parameters():
            param.requires_grad = False

        # 替换最后的分类头
        self.vit.heads = nn.Linear(self.vit.hidden_dim, num_classes)

    def forward(self, x):
        return self.vit(x)


# 3. 训练函数
def train_model(model, trainloader, criterion, optimizer, device, num_epochs=5):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for i, (inputs, labels) in enumerate(trainloader):
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if i % 100 == 99:
                print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 100:.3f}')
                running_loss = 0.0


# 4. 评估函数
def evaluate_model(model, testloader, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in testloader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'Accuracy on test set: {accuracy:.2f}%')
    return accuracy


# 5. 主函数
def main():
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 获取数据
    trainloader, testloader, classes = get_cifar10_data()

    # 创建模型
    model = ViTForCIFAR10(num_classes=10).to(device)

    # 只训练分类头
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.vit.heads.parameters(), lr=0.001)

    # 训练模型
    train_model(model, trainloader, criterion, optimizer, device)

    # 评估模型
    evaluate_model(model, testloader, device)


if __name__ == "__main__":
    main()

# Accuracy on test set: 95.24%
# 对于cifar-100：Final accuracy on test set: 81.11%
